#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

""" This class scale real world parameters to dimension less form
    and call the convection-diffusion-reaction engine to solve the
    parabolic problem."""

from scipy import array, zeros, ones, clip, exp, log

from CDRSolver import CDRSolver

from utils.FemToReal import FemToReal
from utils.Spline import Spline


class GBMSolverAdaptor(CDRSolver):

    """Finite-Element solver for the Black-Scholes PDE

     The problem to solve is given by:

      Function
        f = f(t,x) over the rectangle 0 <= t <= T, Bl <= x <= Bu.
      PDE
        rf = df/dt + rx df/dx + (1/2)(sigma x)^2 d^f/dx^2
      Boundary conditions
        given on the three sides of the rectangle
        t = T; x = Bl; x = Bu

     where r, sigma, T, Bl, Bu are given parameters.
  """

    __slots__ = ['m_spot', 'm_strike', 'm_smin', 'm_smax', 'm_cut_thr',
                 'm_ir', 'm_sigma']

    def __init__(self, time2_mat, spval, strike, ntime_steps, nsegs,
            smin=None, smax=None):
        """Initialize the finite-difference solver.

     Parameters:
      time2_mat - maturity time.
      spval    - Stoke present value.
      strike - strike price.
      ntime_steps - number of time steps to take when discretizing PDE
      nsegs  - number of segments in the finite element discretization
      smin  - lower bound for stock price
      smax  - upper bound for stock price
    """

        self.m_spot = float(spval)
        self.m_strike = float(strike)

        self.m_cut_thr = 1.0e-6
        self.m_smin = float(smin)
        self.m_smax = float(smax)

        if self.m_smin == None or self.m_smax == None:
            self.m_smin = 0.5 * self.m_spot
            self.m_smax = 2.1 * self.m_spot

        (left_point, right_point) = (log(self.m_smin / self.m_strike),
                log(self.m_smax / self.m_strike))
        CDRSolver.__init__(self, nsegs, left_point, right_point)
        (self.m_sigma, self.m_ir) = (None, None)

    def get_strike(self):
        """ get strike value """

        return self.m_strike

    def set_option_cos(self, sigma, dividents, i_rate, s_is_real, q_is_real,
            i_is_real):
        """
      Black scholes is transformed to a diffusion convection reaction
      problem with the following coeffientns

      d = 0.5*self.m_sigma^2
      c = i_rate - dividents - volatility / 2.
      r = -i_rate
    """

        (self.m_sigma, self.m_ir) = (sigma, i_rate)
        if s_is_real:
            diffusion = 0.5 * self.m_sigma ** 2
            convection = self.m_ir - dividents - diffusion
        else:
            assert True == q_is_real and True == i_is_real

            # Extract sigma values, generate new spline using new grid.
            # Change this when introduce adaptivity.

            grid = self.get_mesh()
            yval = array([sigma.value(x) ** 2 for x in grid]) / 2.
            diffusion = Spline(grid.x, yval)
            convection = Spline(grid.x, yval)
            convection.a = (self.m_ir - dividents) \
                * ones(convection.a.shape) - convection.a

        (dl_is_real, sl_is_real) = (s_is_real, s_is_real)
        self.set_diffus_cos(diffusion, convection, -self.m_ir, dl_is_real,
                sl_is_real, i_is_real)

    def _zero(self, pval, time=0):
        """ zero boundary condition. """

        return clip(pval, 0, 0) * clip(time, 0, 0)

    def _ca_vanilla_rbc(self, pval, time=0):
        """ call vanilla right boundary conditions. """

        return self.m_strike * clip(exp(pval) - exp(-self.m_ir * time),
                                    0.0, 1e600)

    def _digital_bc(self, pval, time=0):
        """ digital boundary condition. """

        return exp(-self.m_ir * time) * clip(pval, 1, 1) * clip(time,
                1, 1)

    def _put_va_lbc(self, pval, time=0):
        """ put vanilla left boundary condition. """

        return self.m_strike * clip(exp(-self.m_ir * time) - exp(pval),
                                    0.0, 1e600)

    def _call_va_present_val(self, pval, time=0):
        """ call vanilla present value. """

        return self.m_strike * clip(exp(pval) - 1.0, 0.0, 1e600) \
            * exp(-self.m_ir * time)

    def _call_di_present_val(self, pval, time=0):
        """ digital call present value. """

        tmp = zeros(pval.shape)
        tmp[exp(pval) > 1.0] = exp(-self.m_ir * time)
        return tmp

    def _put_di_present_val(self, pval, time=0):
        """ digital put present value. """

        tmp = zeros(pval.shape)
        tmp[exp(pval) < 1.0] = exp(-self.m_ir * time)
        return tmp

    def _put_va_present_val(self, pval, time=0):
        """ vanilla put present value. """

        return self.m_strike * clip(1.0 - exp(pval), 0.0, 1e600) \
            * exp(-self.m_ir * time)

    def _post_process(self, u_final):
        """ finite element to real space. """

        grid = self.get_mesh()
        basis = self.get_sfunc()

        recons = FemToReal(grid, basis)
        solu = recons.reconstruct(u_final)
        solu[:, 0] = self.m_strike * exp(solu[:, 0])  # S
        price = solu[:, 1]
        price[price < self.m_cut_thr] = 0.0
        return solu

    def get_stock_pres_val(self):
        """ stock present value. """

        return self.m_spot

    def get_price_history(self):
        """ price history. """

        price_hist = None
        if self.get_keep_hist():
            grid = self.get_mesh()
            recons = FemToReal(grid, self.get_sfunc())
            from scipy import arange

            alg_hist = self._get_solver_hist()
            ntime_steps = alg_hist.shape[0]
            price_hist = zeros((ntime_steps, 10 * grid.get_nsegs()),
                               'float')
            for i in arange(ntime_steps):
                solu = recons.reconstruct(alg_hist[i, :])
                price = solu[:, 1]
                price[price < self.m_cut_thr] = 0.0
                price_hist[i, :] = price.flatten()
        return price_hist


